package cn.iocoder.yudao.module.engine.service.core.messagestore;

import cn.iocoder.yudao.module.engine.anno.Syn;
import cn.iocoder.yudao.module.engine.dal.redis.engine.EngineLockRedisDAO;
import cn.iocoder.yudao.module.engine.model.config.EngineConfig;
import cn.iocoder.yudao.module.engine.model.context.AIMessageContext;
import cn.iocoder.yudao.module.engine.model.context.ControlEngineContext;
import cn.iocoder.yudao.module.engine.model.engine.EngineOutputData;
import cn.iocoder.yudao.module.infra.utils.EngineExecutorUtils;
import cn.iocoder.yudao.module.infra.util.GameMatchIdUtil;
import cn.iocoder.yudao.module.infra.util.TraceIdUtil;
import jakarta.annotation.PostConstruct;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;

import javax.annotation.Nullable;

import jakarta.annotation.Resource;

import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.concurrent.ExecutorService;

@Order(100)
@Component
public class MessageStoreManager {
    private ExecutorService saveExecutors;

    private final Map<String, MessageStore> storeMap = Collections.synchronizedMap(new LinkedHashMap<>());

    @Resource
    private EngineConfig engineConfig;

    @Resource
    private EngineLockRedisDAO engineLockRedisDAO;

    @PostConstruct
    public void init() {
        final EngineConfig.ThreadPool threadPool = engineConfig.getThreadPool();
        saveExecutors = EngineExecutorUtils.newThreadPoolExecutor(
                threadPool.getMessageSaveCorePoolSize(),
                "message-save"
        );
    }

    public void registerStore(String name, MessageStore store) {
        storeMap.put(name, store);
    }

    public MessageStore getStore(String name) {
        return storeMap.getOrDefault(name, null);
    }

    public void saveMessages(ControlEngineContext engineContext,
                             @Nullable AIMessageContext messageCtx,
                             EngineOutputData.Audio audio) {
        // syn
        storeMap.forEach((name, store) -> {
            final Syn synAnnotation = store.getClass().getAnnotation(Syn.class);
            if (synAnnotation != null) {
                store.saveMessages(engineContext, messageCtx, audio);
            }
        });

        if (engineConfig.getMock()) {
            return;
        }

        // async
        final String messageId = audio.getChatId();
        final String traceId = TraceIdUtil.getTraceId();
        saveExecutors.execute(() -> {
            try {
                TraceIdUtil.setUserId(engineContext.getUserId());
                TraceIdUtil.setTraceId(traceId);
                GameMatchIdUtil.setMatchId(engineContext.getGameMatchId());
                engineLockRedisDAO.lockMessage(messageId, () -> {
                    for (MessageStore store : storeMap.values()) {
                        final Syn synAnnotation = store.getClass().getAnnotation(Syn.class);
                        if (synAnnotation == null) {
                            store.saveMessages(engineContext, messageCtx, audio);
                        }
                    }
                    return null;
                });
            } finally {
                TraceIdUtil.removeUserId();
                TraceIdUtil.removeTraceId();
                GameMatchIdUtil.removeMatchId();
            }
        });
    }

}
